package goleetcode

import (
	"strconv"
)

func punishmentNumber(n int) int {
	res := 0
	for i := 1; i <= n; i++ {
		if dfs(strconv.Itoa(i*i), 0, 0, i) {
			res += i * i
		}
	}
	return res
}

func dfs(s string, pos int, tot int, target int) bool {
	if pos == len(s) {
		return tot == target
	}

	var sum int = 0
	for i := pos; i < len(s); i++ {
		sum = sum*10 + int(s[i]-'0')
		if sum+tot > target {
			break
		}
		if dfs(s, i+1, sum+tot, target) {
			return true
		}
	}
	return false
}
